"""
Star-shaped interface experiment for AE-PINN validation.

This experiment implements Example 4.3 from paper 2506.18332, featuring a star-shaped
interface defined by the level-set function from Equation 3.17:

φ(x,y) = √(x² + y²) - r₀(1 + Σ βₖ cos(ηₖ(arctan(y/x) - θₖ)))

The experiment demonstrates the AE-PINN's ability to handle complex interface geometries
with prescribed jump conditions across the interface.
"""

import torch
import torch.nn as nn
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
import matplotlib.pyplot as plt
import os
import sys
from pathlib import Path

# Add parent directory to path for imports
sys.path.append(str(Path(__file__).parent.parent))

from networks.fcnn import FCNN
from networks.ia_nn import IANN, CompositeAEPINN, create_star_shaped_level_set
from pinn.loss import AEPINNLoss
from pinn.solver import AEPINNSolver


class StarShapedInterfaceProblem:
    """
    Star-shaped interface problem setup.
    
    Defines the complete elliptic interface problem with:
    - Star-shaped interface geometry
    - Piecewise constant diffusion coefficient
    - Known analytical solution for validation
    - Prescribed jump conditions
    """
    
    def __init__(self, r0=0.3, alpha_in=1.0, alpha_out=5.0, 
                 betas=None, etas=None, thetas=None):
        """
        Initialize star-shaped interface problem.
        
        Args:
            r0 (float): Base radius of star-shaped interface
            alpha_in (float): Diffusion coefficient inside interface
            alpha_out (float): Diffusion coefficient outside interface
            betas (list): Star shape amplitude coefficients
            etas (list): Star shape frequency coefficients
            thetas (list): Star shape phase coefficients
        """
        self.r0 = r0
        self.alpha_in = alpha_in
        self.alpha_out = alpha_out
        
        # Star shape parameters
        if betas is None:
            self.betas = [0.3, 0.2, 0.1]
        else:
            self.betas = betas
            
        if etas is None:
            self.etas = [3, 4, 5]
        else:
            self.etas = etas
            
        if thetas is None:
            self.thetas = [0, np.pi/4, np.pi/2]
        else:
            self.thetas = thetas
        
        # Domain bounds
        self.domain_bounds = [[-1.0, 1.0], [-1.0, 1.0]]
        
        # Create level-set function
        self.level_set_func = create_star_shaped_level_set(
            r0=self.r0, betas=self.betas, etas=self.etas, thetas=self.thetas
        )
        
        # Problem functions
        self.alpha_func = self._create_alpha_function()
        self.f_func = self._create_source_function()
        self.boundary_func = self._create_boundary_function()
        self.jump_func = self._create_jump_function()
        self.flux_jump_func = self._create_flux_jump_function()
        
        # Analytical solution (if available)
        self.analytical_func = self._create_analytical_solution()
    
    def _create_alpha_function(self):
        """Create piecewise constant diffusion coefficient."""
        def alpha_func(x, y):
            phi = self.level_set_func(x, y)
            return torch.where(phi < 0, self.alpha_in, self.alpha_out)
        return alpha_func
    
    def _create_source_function(self):
        """Create source term f(x,y)."""
        def f_func(x, y):
            # For this example, use a simple source term
            return torch.sin(np.pi * x) * torch.sin(np.pi * y)
        return f_func
    
    def _create_boundary_function(self):
        """Create boundary condition function."""
        def boundary_func(x, y):
            # Homogeneous Dirichlet boundary conditions
            return torch.zeros_like(x)
        return boundary_func
    
    def _create_jump_function(self):
        """Create interface jump condition function."""
        def jump_func(x, y):
            # Prescribed jump [[u]] = β(x,y)
            # For this example, use a simple jump condition
            return 0.5 * torch.cos(2 * np.pi * (x + y))
        return jump_func
    
    def _create_flux_jump_function(self):
        """Create interface flux jump condition function."""
        def flux_jump_func(x, y):
            # Prescribed flux jump [[α∇u]]·n = ρ(x,y)
            # For this example, use zero flux jump
            return torch.zeros_like(x)
        return flux_jump_func
    
    def _create_analytical_solution(self):
        """Create analytical solution for validation (if known)."""
        # For this complex geometry, analytical solution is not easily available
        # We use numerical reference solution or manufactured solution
        def analytical_func(x, y):
            # Placeholder: could be implemented with reference numerical solution
            return torch.zeros_like(x)
        return analytical_func
    
    def plot_interface_geometry(self, save_path=None):
        """Plot the star-shaped interface geometry."""
        x = torch.linspace(-1, 1, 400)
        y = torch.linspace(-1, 1, 400)
        X, Y = torch.meshgrid(x, y, indexing='ij')
        
        # Compute level-set function
        phi = self.level_set_func(X.flatten(), Y.flatten()).reshape(X.shape)
        
        # Create plot
        fig, ax = plt.subplots(1, 1, figsize=(8, 8))
        
        # Plot interface (φ = 0 contour)
        contour = ax.contour(X.numpy(), Y.numpy(), phi.numpy(), levels=[0], colors='red', linewidths=3)
        ax.clabel(contour, inline=True, fontsize=12, fmt='Interface')
        
        # Plot level-set function
        im = ax.contourf(X.numpy(), Y.numpy(), phi.numpy(), levels=20, alpha=0.6, cmap='RdBu')
        plt.colorbar(im, ax=ax, label='Level-set φ(x,y)')
        
        # Formatting
        ax.set_xlabel('x')
        ax.set_ylabel('y')
        ax.set_title('Star-Shaped Interface Geometry')
        ax.set_aspect('equal')
        ax.grid(True, alpha=0.3)
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Interface geometry saved to {save_path}")
        else:
            plt.savefig('interface_geometry.png', dpi=300, bbox_inches='tight')
            print("Interface geometry saved to interface_geometry.png")
        plt.close()
    
    def plot_diffusion_coefficient(self, save_path=None):
        """Plot the piecewise diffusion coefficient."""
        x = torch.linspace(-1, 1, 200)
        y = torch.linspace(-1, 1, 200)
        X, Y = torch.meshgrid(x, y, indexing='ij')
        
        # Compute diffusion coefficient
        alpha_vals = self.alpha_func(X.flatten(), Y.flatten()).reshape(X.shape)
        
        # Create plot
        fig, ax = plt.subplots(1, 1, figsize=(8, 8))
        
        # Plot diffusion coefficient
        im = ax.contourf(X.numpy(), Y.numpy(), alpha_vals.numpy(), 
                        levels=[self.alpha_in, (self.alpha_in + self.alpha_out)/2, self.alpha_out],
                        colors=['blue', 'red'], alpha=0.7)
        plt.colorbar(im, ax=ax, label='Diffusion coefficient α(x,y)')
        
        # Plot interface
        phi = self.level_set_func(X.flatten(), Y.flatten()).reshape(X.shape)
        ax.contour(X.numpy(), Y.numpy(), phi.numpy(), levels=[0], colors='black', linewidths=2)
        
        # Formatting
        ax.set_xlabel('x')
        ax.set_ylabel('y')
        ax.set_title(f'Piecewise Diffusion Coefficient (α_in={self.alpha_in}, α_out={self.alpha_out})')
        ax.set_aspect('equal')
        ax.grid(True, alpha=0.3)
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Interface geometry saved to {save_path}")
        else:
            plt.savefig('interface_geometry.png', dpi=300, bbox_inches='tight')
            print("Interface geometry saved to interface_geometry.png")
        plt.close()


class StarShapedExperiment:
    """
    Complete experimental setup for star-shaped interface problem.
    
    Handles model creation, training, and result analysis for the
    star-shaped interface validation experiment.
    """
    
    def __init__(self, problem_config=None, model_config=None, training_config=None):
        """
        Initialize experiment.
        
        Args:
            problem_config (dict): Problem configuration parameters
            model_config (dict): Model architecture configuration
            training_config (dict): Training configuration parameters
        """
        # Default configurations
        self.problem_config = problem_config or {
            'r0': 0.3,
            'alpha_in': 1.0,
            'alpha_out': 5.0,
            'betas': [0.3, 0.2, 0.1],
            'etas': [3, 4, 5],
            'thetas': [0, np.pi/4, np.pi/2]
        }
        
        self.model_config = model_config or {
            'fcnn_layers': [50, 50, 50],
            'ia_hidden_dim': 40,
            'ia_attention_layers': 3,
            'activation': 'tanh'
        }
        
        self.training_config = training_config or {
            'epochs': 15000,
            'learning_rate': 1e-3,
            'n_interior': 2000,
            'n_boundary': 400,
            'n_interface': 400,
            'optimizer': 'adam',
            'scheduler': 'step',
            'loss_weights': {
                'tau_pde': 1.0,
                'tau_bc': 100.0,
                'tau_interface_u': 100.0,
                'tau_interface_flux': 100.0
            }
        }
        
        # Create problem
        self.problem = StarShapedInterfaceProblem(**self.problem_config)
        
        # Initialize model and solver
        self.model = None
        self.solver = None
    
    def create_model(self):
        """Create the composite AE-PINN model."""
        # FCNN for continuous component
        fcnn = FCNN(
            input_dim=2,
            output_dim=1,
            hidden_layers=self.model_config['fcnn_layers'],
            activation=self.model_config['activation']
        )
        
        # IA-NN for discontinuous component
        ia_nn = IANN(
            input_dim=2,
            hidden_dim=self.model_config['ia_hidden_dim'],
            num_attention_layers=self.model_config['ia_attention_layers'],
            output_dim=1,
            level_set_func=self.problem.level_set_func,
            activation=self.model_config['activation']
        )
        
        # Composite model
        self.model = CompositeAEPINN(fcnn, ia_nn)
        
        print(f"Created AE-PINN model with {sum(p.numel() for p in self.model.parameters())} parameters")
        print(f"FCNN parameters: {sum(p.numel() for p in fcnn.parameters())}")
        print(f"IA-NN parameters: {sum(p.numel() for p in ia_nn.parameters())}")
        
        return self.model
    
    def create_solver(self):
        """Create the AE-PINN solver."""
        if self.model is None:
            self.create_model()
        
        # Create loss function
        loss_weights = self.training_config['loss_weights']
        loss_function = AEPINNLoss(
            alpha_func=self.problem.alpha_func,
            f_func=self.problem.f_func,
            level_set_func=self.problem.level_set_func,
            tau_pde=loss_weights['tau_pde'],
            tau_bc=loss_weights['tau_bc'],
            tau_interface_u=loss_weights['tau_interface_u'],
            tau_interface_flux=loss_weights['tau_interface_flux']
        )
        
        # Create solver
        self.solver = AEPINNSolver(
            model=self.model,
            loss_function=loss_function,
            domain_bounds=self.problem.domain_bounds,
            level_set_func=self.problem.level_set_func
        )
        
        # Setup optimizer
        self.solver.setup_optimizer(
            optimizer_type=self.training_config['optimizer'],
            learning_rate=self.training_config['learning_rate'],
            scheduler_type=self.training_config['scheduler']
        )
        
        return self.solver
    
    def run_training(self, save_path=None):
        """Run the complete training process."""
        if self.solver is None:
            self.create_solver()
        
        print("Starting star-shaped interface experiment training...")
        
        # Run training
        self.solver.train(
            epochs=self.training_config['epochs'],
            n_interior=self.training_config['n_interior'],
            n_boundary=self.training_config['n_boundary'],
            n_interface=self.training_config['n_interface'],
            boundary_func=self.problem.boundary_func,
            jump_func=self.problem.jump_func,
            flux_jump_func=self.problem.flux_jump_func,
            print_every=1000,
            save_every=5000,
            save_path=save_path,
            use_adaptive_weights=True
        )
        
        print("Training completed!")
    
    def evaluate_solution(self, n_test=10000):
        """Evaluate the trained solution on test points."""
        if self.model is None:
            raise ValueError("Model not trained yet!")
        
        self.model.eval()
        
        # Initialize CUDA context if using GPU to avoid cuBLAS context issues
        device = next(self.model.parameters()).device
        if device.type == 'cuda':
            torch.cuda.empty_cache()
            # Perform a simple operation to initialize CUDA context
            dummy_tensor = torch.zeros(1, device=device)
            _ = dummy_tensor + 1
            torch.cuda.synchronize()
        
        # Generate test points
        x_test = torch.rand(n_test, 2) * 2 - 1  # Uniform in [-1, 1]²
        x_test = x_test.to(next(self.model.parameters()).device)
        
        with torch.no_grad():
            u_pred = self.model(x_test)
        
        return x_test.cpu(), u_pred.cpu()
    
    def plot_solution(self, save_path=None, resolution=200):
        """Plot the predicted solution."""
        if self.model is None:
            raise ValueError("Model not trained yet!")
        
        self.model.eval()
        
        # Initialize CUDA context if using GPU to avoid cuBLAS context issues
        device = next(self.model.parameters()).device
        if device.type == 'cuda':
            torch.cuda.empty_cache()
            # Perform a simple operation to initialize CUDA context
            dummy_tensor = torch.zeros(1, device=device)
            _ = dummy_tensor + 1
            torch.cuda.synchronize()
        
        # Create mesh for plotting
        x = torch.linspace(-1, 1, resolution)
        y = torch.linspace(-1, 1, resolution)
        X, Y = torch.meshgrid(x, y, indexing='ij')
        
        points = torch.stack([X.flatten(), Y.flatten()], dim=1)
        points = points.to(next(self.model.parameters()).device)
        
        # Predict solution
        with torch.no_grad():
            u_pred = self.model(points).cpu().numpy().reshape(X.shape)
        
        # Compute level-set for interface
        phi = self.problem.level_set_func(X.flatten(), Y.flatten()).reshape(X.shape)
        
        # Create plot
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 7))
        
        # Plot solution
        im1 = ax1.contourf(X.numpy(), Y.numpy(), u_pred, levels=20, cmap='viridis')
        ax1.contour(X.numpy(), Y.numpy(), phi.numpy(), levels=[0], colors='red', linewidths=2)
        plt.colorbar(im1, ax=ax1, label='Solution u(x,y)')
        ax1.set_title('AE-PINN Solution')
        ax1.set_xlabel('x')
        ax1.set_ylabel('y')
        ax1.set_aspect('equal')
        
        # Plot solution with interface highlighted
        im2 = ax2.contourf(X.numpy(), Y.numpy(), u_pred, levels=20, cmap='RdBu_r')
        ax2.contour(X.numpy(), Y.numpy(), phi.numpy(), levels=[0], colors='black', linewidths=3)
        ax2.contourf(X.numpy(), Y.numpy(), phi.numpy(), levels=[0, 1], colors=['none', 'gray'], alpha=0.3)
        plt.colorbar(im2, ax=ax2, label='Solution u(x,y)')
        ax2.set_title('Solution with Interface Regions')
        ax2.set_xlabel('x')
        ax2.set_ylabel('y')
        ax2.set_aspect('equal')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Interface geometry saved to {save_path}")
        else:
            plt.savefig('interface_geometry.png', dpi=300, bbox_inches='tight')
            print("Interface geometry saved to interface_geometry.png")
        plt.close()
    
    def run_complete_experiment(self, save_dir=None):
        """Run the complete experimental pipeline."""
        if save_dir:
            os.makedirs(save_dir, exist_ok=True)
        
        print("="*60)
        print("AE-PINN Star-Shaped Interface Experiment")
        print("="*60)
        
        # Plot problem setup
        print("1. Plotting problem geometry...")
        self.problem.plot_interface_geometry(
            save_path=os.path.join(save_dir, 'interface_geometry.png') if save_dir else None
        )
        self.problem.plot_diffusion_coefficient(
            save_path=os.path.join(save_dir, 'diffusion_coefficient.png') if save_dir else None
        )
        
        # Create and train model
        print("2. Creating and training AE-PINN model...")
        self.create_solver()
        self.run_training(save_path=save_dir)
        
        # Plot training history
        print("3. Analyzing training results...")
        self.solver.plot_training_history(
            save_path=os.path.join(save_dir, 'training_history.png') if save_dir else None
        )
        
        # Evaluate and plot solution
        print("4. Evaluating and plotting solution...")
        self.plot_solution(
            save_path=os.path.join(save_dir, 'solution.png') if save_dir else None
        )
        
        # Save final model
        if save_dir:
            torch.save(self.model.state_dict(), os.path.join(save_dir, 'final_model.pt'))
        
        print("Experiment completed successfully!")


if __name__ == "__main__":
    # Run star-shaped interface experiment
    experiment = StarShapedExperiment()
    
    # Run complete experiment
    save_directory = "/home/linden/Downloads/arxiv_paper/RINN/results/star_shaped_experiment"
    experiment.run_complete_experiment(save_dir=save_directory)